Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@bionemo-recipes/models/esm2/modeling_esm_te.py`:
- Around line 214-234: Add unit tests that exercise the per-layer FP context
selection in the loop that iterates self.layers and reads
self.layer_number_quantized_recipe_map: create three scenarios where the mapped
fp_recipe for a given layer is (a) an instance of FP8_RECIPES, (b) an instance
of FP4_RECIPES, and (c) None/other; for each case assert that fp_context becomes
nullcontext() for FP8, transformer_engine.pytorch.autocast(enabled=True,
recipe=fp_recipe) for FP4, and
transformer_engine.pytorch.autocast(enabled=False) for the default/BF16 path.
Use lightweight mocks/monkeypatching to replace
transformer_engine.pytorch.autocast with a stub that records its args/returns so
you can assert enabled and recipe values, and construct minimal model instances
(or unit-test the loop function directly) that set
layer_number_quantized_recipe_map and self.layers to trigger each branch; also
include a test that output_hidden_states True appends hidden_states to
all_hidden_states.

In `@bionemo-recipes/recipes/esm2_native_te/example_8m_checkpoint/esm_nv.py`:
- Line 234: The TODO on the forward path flags unverified precision-routing
behavior; replace it by adding focused unit tests that exercise the precision
routing logic for FP8, FP4 and BF16 (e.g., create tests that run the module's
forward pass with tensors/configs that should route through each precision
branch and assert the correct branch was used and outputs match expected
numerical/shape properties), and then remove the TODO comment. Target the
functions and code paths that implement precision routing (the forward path /
precision-routing conditional blocks referenced by the TODO in esm_nv.py) and
add parametrized tests that cover boundary cases and mixed precision
combinations to ensure deterministic routing.

In `@bionemo-recipes/recipes/esm2_native_te/fp4_debugging_stats.yaml`:
- Around line 21-23: Update the comment describing layer ranges so it matches
the actual regex stored in layer_name_regex_pattern; the current comment
mentions layers 0-4 / 1-5 but the regex
'model\.esm\.encoder\.layers\.([6-9]|10)\..*(layernorm_qkv|proj|fc1|fc2)'
targets layers 6–10, so change the comment to state layers 6-10 (or 7-11 if
using 1-indexed wording) to prevent confusion and ensure the comment and the
pattern are consistent.

In `@bionemo-recipes/recipes/esm2_native_te/fp8_debugging_stats.yaml`:
- Around line 19-23: Replace the invalid tensor type "fprop" in the
LogTensorStats block with the valid TransformerEngine tensor type "activation"
(modify the tensors: [dgrad, wgrad, fprop] entry to tensors: [dgrad, wgrad,
activation]); keep stats and freq as-is, or optionally refactor LogTensorStats
to use tensors_struct for per-tensor configs if you need different stats per
tensor type.

In `@bionemo-recipes/recipes/esm2_native_te/hydra_config/defaults.yaml`:
- Around line 55-61: Either remove the unused fp4_model_init_kwargs entry from
the fp4_config block in defaults.yaml, or implement FP4 model initialization in
train_fsdp2.py to mirror the FP8 pattern: when args.fp4_config.enabled and
args.fp4_config.fp4_model_init_kwargs.enabled are true, call
transformer_engine.pytorch.quantized_model_init(...) (or the appropriate FP4
init API) with args.fp4_config.fp4_model_init_kwargs before model training;
update any related code paths that reference
fp4_config.fp4_recipe/fp4_format/fp4_recipe_kwargs to ensure consistent
behavior.

In `@bionemo-recipes/recipes/esm2_native_te/modeling_esm_te.py`:
- Around line 130-133: The assertion message for the padded_vocab_size check is
too long; update the assert in the block that checks self.padded_vocab_size and
self.vocab_size to keep lines <=119 characters by moving the long f-string into
a shorter expression (e.g., build the message in a local variable like msg =
(f"padded_vocab_size ({self.padded_vocab_size}) must be greater than or equal to
" f"vocab_size ({self.vocab_size})") or use implicit string concatenation/split
across lines) and then call assert self.padded_vocab_size >= self.vocab_size,
msg; keep references to self.padded_vocab_size and self.vocab_size so the
semantic check and error content remain unchanged.

In `@bionemo-recipes/recipes/esm2_native_te/quantization.py`:
- Around line 86-94: The temp YAML file created with "temp_file =
tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False)" is left on
disk because you return temp_file.name; change this by either (A) making the
function return the config contents or a context-managed path (use
tempfile.NamedTemporaryFile with delete=True or write via
tempfile.TemporaryDirectory and yield the path) so the file is removed
automatically, or (B) accept a caller-provided "quant_log_dir" or "cleanup" flag
and write the file into that deterministic log directory (or delete the temp
file before returning when appropriate). Update the function's docstring and any
callers to reflect whether the caller is responsible for cleanup, and adjust
logging (logger.info uses temp_file.name) to log the deterministic path or note
that the file is ephemeral. Ensure you modify the code around the
NamedTemporaryFile usage and the return value accordingly.

In `@bionemo-recipes/recipes/esm2_native_te/README.md`:
- Line 106: Replace the vague anchor text "here" with a descriptive label that
explains the target, e.g., change "and
[here](https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/examples/fp8_primer.html)"
to something like "and the NVIDIA Transformer Engine FP8 primer" (update the
anchor text around the existing URL), so the README sentence reads clearly and
satisfies markdown linting.
- Around line 80-83: Fix the wording in the low-precision benchmark paragraph:
change "low precision" to "low-precision" throughout, add missing period after
"etc" ("etc."), correct the typo "outweights" to "outweighs", and rephrase the
sentence that reads "the cost to quantize activations from high precision to low
precision outweights the benefits of performing matrix multiplication with low
precision" to a clearer form (e.g., "the cost to quantize activations from
high-precision to low-precision outweighs the benefits of using low-precision
matrix multiplication") so the paragraph reads smoothly and consistently.

In `@bionemo-recipes/recipes/esm2_native_te/train_ddp.py`:
- Around line 62-64: Replace the silent warning when args.fp4_config.enabled is
true with a fail-fast guard: in the block that currently calls logger.warning
(the check of args.fp4_config.enabled), either raise a clear RuntimeError to
stop execution or require an explicit override flag (e.g.,
args.allow_experimental_nvfp4_ddp) before proceeding; update the check to
validate that if fp4 is enabled and DDP is in use (the current DDP launch path),
then if not args.allow_experimental_nvfp4_ddp raise an error with a message
explaining NVFP4+DDP is unsupported and how to opt into experimental mode.
Ensure you reference and change the condition around args.fp4_config.enabled and
the logger.warning call so the run is blocked unless the explicit override is
provided.

In `@bionemo-recipes/recipes/esm2_peft_te/example_8m_checkpoint/esm_nv.py`:
- Around line 214-234: Add unit tests validating the per-layer quantization
context selection in the loop that uses layer_number_quantized_recipe_map:
ensure that when layer_number_quantized_recipe_map returns an FP8 recipe the
code sets fp_context to nullcontext(), when it returns an FP4 recipe it uses
transformer_engine.pytorch.autocast(enabled=True, recipe=fp_recipe), and for
None/other recipes it uses transformer_engine.pytorch.autocast(enabled=False);
add tests covering output_hidden_states path as well (all_hidden_states
behavior) and include edge cases (missing map, unexpected recipe types), and
also correct the typo in the TODO comment from "funciton" to "function".

---

Nitpick comments:
In `@bionemo-recipes/models/esm2/modeling_esm_te.py`:
- Around line 60-67: FP4_RECIPES is defined as a single value whereas
FP8_RECIPES is a tuple; make FP4_RECIPES a tuple for consistency by wrapping
transformer_engine.common.recipe.NVFP4BlockScaling in a one-element tuple (use a
trailing comma) so the constant mirrors FP8_RECIPES and any tuple-based handling
of recipes (reference FP4_RECIPES and FP8_RECIPES names).

In `@bionemo-recipes/recipes/esm2_native_te/.dockerignore`:
- Around line 33-34: The ignore pattern "j/" in .dockerignore matches any
directory named "j" at any depth; change it to "/j/" to anchor it to the
repository root so only the top-level scratch dir is ignored. Update the pattern
"j/" to "/j/" (preserve the trailing slash) in the .dockerignore entry to avoid
accidentally excluding nested directories named "j".

In `@bionemo-recipes/recipes/esm2_native_te/dataset.py`:
- Around line 60-62: The call to AutoTokenizer.from_pretrained uses a redundant
conditional for the revision argument; replace revision=tokenizer_revision if
tokenizer_revision else None with simply revision=tokenizer_revision in the
AutoTokenizer.from_pretrained(...) call so the revision parameter directly uses
the tokenizer_revision variable.

In `@bionemo-recipes/recipes/esm2_native_te/hydra_config/defaults.yaml`:
- Around line 104-107: The default for use_fp32_master_weights is set to null
which may be unexpected when later checked as args.use_fp32_master_weights;
update the defaults.yaml to document that null is intentional and treated as
falsy (or change the default to false) and add a short comment next to
use_fp32_master_weights explaining that the training script expects a
boolean-like value and that null will be treated as false by the conditional
using args.use_fp32_master_weights; ensure any consumers (e.g., the code that
checks args.use_fp32_master_weights) handle null explicitly if you want
different behavior.

In `@bionemo-recipes/recipes/esm2_native_te/hydra_config/L1_650M.yaml`:
- Around line 9-12: Add a blank line between the dataset block and the WandB
section to match other config files: locate the dataset section (keys dataset,
micro_batch_size, tokenizer_revision) and insert a single empty line before the
existing WandB comment/section (the "# WandB config" or the wandb_init_args
section) so formatting is consistent with L1_3B.yaml and L1_15B_perf_test.yaml.

In `@bionemo-recipes/recipes/esm2_native_te/modeling_esm_te.py`:
- Around line 604-609: Replace the hardcoded mask ratio literal by a named class
constant and use it wherever needed: define a class-level constant (e.g.,
MASK_RATIO_TRAIN = 0.15 * 0.8) on the model class, then update the local
variable usage in the method that currently sets mask_ratio_train and in the
similar block at lines 625-637 to read from that constant (self.MASK_RATIO_TRAIN
or ClassName.MASK_RATIO_TRAIN) so scale_factor and embedding scaling use the
named constant for clarity and potential configurability.
- Around line 228-234: Extract the per-layer FP context selection into a small
function (e.g., get_fp_context(fp_recipe)) that returns nullcontext when
fp_recipe is in FP8_RECIPES, returns
transformer_engine.pytorch.autocast(enabled=True, recipe=fp_recipe) when
fp_recipe is in FP4_RECIPES, and returns
transformer_engine.pytorch.autocast(enabled=False) otherwise; then add a unit
test file that parametrizes several fp_recipe values (members of FP8_RECIPES,
FP4_RECIPES, and a default/None case) and asserts that get_fp_context returns
the expected context object type/behavior (e.g., is nullcontext for FP8 cases
and an autocast context for FP4/default) to validate the per-layer FP context
selection logic.

In `@bionemo-recipes/recipes/esm2_native_te/quantization.py`:
- Around line 61-62: The file open call using config_file should specify an
explicit encoding to avoid platform-dependent behavior; update the open(...) in
the block that reads the YAML (the with open(config_file, "r") as f: config =
yaml.safe_load(f) statement) to include encoding="utf-8" so the YAML is read
consistently across platforms.
- Around line 136-157: Replace the manual boilerplate class QuantizationLayers
with a dataclass: add "from dataclasses import dataclass" and annotate the class
with `@dataclass`, convert the four constructor args (fp8_layers_0indexed,
fp4_layers_0indexed, fp8_layers_1indexed, fp4_layers_1indexed) to dataclass
fields using Optional[list[int]] types (or list[int] | None) and remove the
explicit __init__; keep the existing class docstring and attribute names so
__repr__ and __eq__ are provided automatically.

In `@bionemo-recipes/recipes/esm2_native_te/tests/test_train.py`:
- Around line 146-158: Rename test identifiers and docstring to reflect the
config rename from fp8_stats_config to quant_stats_config: update the test
function name test_sanity_ddp_fp8_stats_logging to
test_sanity_ddp_quant_stats_logging, change the fp8_log_dir variable to
quant_log_dir (or similar), and update the docstring "FP8 stats logging" to
"quant stats logging" while keeping all uses of quant_stats_config and the
existing assertions intact so the test still validates FP8 behavior under the
new config name.

In `@bionemo-recipes/recipes/esm2_peft_te/example_8m_checkpoint/esm_nv.py`:
- Around line 60-67: FP4_RECIPES is assigned a single class
(transformer_engine.common.recipe.NVFP4BlockScaling) while FP8_RECIPES is a
tuple; make FP4_RECIPES a tuple for consistency and to allow adding more entries
later — update the FP4_RECIPES assignment to use a tuple containing
NVFP4BlockScaling (e.g., FP4_RECIPES =
(transformer_engine.common.recipe.NVFP4BlockScaling,)) so code that expects a
sequence of recipe classes (similar to FP8_RECIPES) will work uniformly.

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@bionemo-recipes/models/esm2/tests/test_layer_quantization.py`:
- Around line 16-242: Add a golden-value parity test that runs the
TransformerEngine ESM encoder (NVEsmEncoder) and the reference ESM encoder on
the same deterministic input/seed and asserts numerical parity (e.g., final
token logits or pooled embeddings) within a small tolerance; create a new test
function (e.g., test_te_vs_reference_golden_value_parity) in this module that
uses torch.manual_seed, a small random input tensor on CUDA, constructs an
NVEsmEncoder via NVEsmConfig and constructs the reference ESM model (import the
reference model used in the repo), runs both forward passes with identical
settings, and asserts outputs are close with pytest.approx or torch.allclose;
ensure the test uses the existing encoder fixture pattern/device and keeps the
comparison deterministic and tolerant to tiny numeric differences.

In `@bionemo-recipes/recipes/esm2_native_te/modeling_esm_te.py`:
- Around line 213-222: In initialize_quantization, validate fp8_layers and
fp4_layers before applying them: convert to sets (fp8_layers_set,
fp4_layers_set), check every layer id is an int within 0..len(self.layers)-1 and
raise a ValueError if any id is out of range, check for overlap by computing
intersection = fp8_layers_set & fp4_layers_set and raise a ValueError if
non-empty, and optionally ensure inputs are unique/convertible to int; only
after these checks populate self._layer_precision using range(len(self.layers)).

In `@bionemo-recipes/recipes/esm2_native_te/train_ddp_cp.py`:
- Around line 78-81: The FP8 recipe is being constructed unconditionally
(fp8_recipe via hydra.utils.get_class and Format[args.fp8_config.fp8_format])
even when FP8 is disabled; wrap that construction in an if
args.fp8_config.enabled: guard (same pattern used for fp4_recipe) so fp8_recipe
is only created when args.fp8_config.enabled is true, and ensure any references
to fp8_recipe are only used within that guarded block or handled when disabled.

---

Nitpick comments:
In `@bionemo-recipes/recipes/esm2_native_te/train_fsdp2_cp.py`:
- Around line 87-90: The FP8 recipe is created unconditionally which forces FP8
recipe validation even when FP8 is disabled; wrap the hydra.utils.get_class(...)
call that constructs fp8_recipe in a conditional that checks
args.fp8_config.enabled and only builds fp8_recipe (using
Format[args.fp8_config.fp8_format] and args.fp8_config.fp8_recipe_kwargs) when
enabled, otherwise set fp8_recipe to None or skip creation so BF16/FP4-only runs
don't require FP8 recipe validity.

In `@bionemo-recipes/recipes/esm2_native_te/train_mfsdp.py`:
- Around line 69-77: fp8_recipe is always instantiated even when FP8 is
disabled; change the logic to only construct fp8_recipe when
args.fp8_config.enabled is true (mirror how fp4_recipe is handled).
Specifically, initialize fp8_recipe = None and wrap the
hydra.utils.get_class(...) construction in an if args.fp8_config.enabled: block
so that fp8_recipe is only created when enabled, referencing the existing
symbols fp8_recipe and args.fp8_config.enabled to locate and modify the code.